from langchain.agents import initialize_agent
from langchain.tools import BaseTool
from langchain_ollama import OllamaLLM

from common_config import CHAT_OLLAMA_MODEL


# 定义一个自定义工具：计算平方和
class SquareSumTool(BaseTool):
    name: str = "square_sum_tool"  # 添加类型注解
    description: str = "A tool_call to calculate the sum of squares of a list of numbers."

    def _run(self, numbers: str):
        print(f"Numbers: {numbers}")
        # 将输入的字符串转换为数字列表
        number_list = [int(num) for num in numbers.split(",")]
        # 计算平方和
        result = sum([num ** 2 for num in number_list])
        print(f"Result: {result}")
        return str(result)

    async def _arun(self, numbers: str):
        raise NotImplementedError("This tool_call does not support async.")

# 初始化Ollama模型
ollama = OllamaLLM(model=CHAT_OLLAMA_MODEL)

# 初始化LangChain代理，传入工具
tools = [SquareSumTool()]
agent = initialize_agent(tools, ollama, verbose=True)

# 调用代理，使用工具
input_query = """
Calculate the sum of squares for the numbers 1, 2, 3 and 4.

/nothink
"""
response = agent.run(input_query)

print(f"Response: {response}")